package top.ochiamalu.aioj.util;

import com.google.gson.Gson;
import com.google.gson.JsonArray;
import com.google.gson.JsonObject;
import com.google.gson.JsonSyntaxException;
import lombok.extern.slf4j.Slf4j;
import okhttp3.Call;
import okhttp3.Callback;
import okhttp3.MediaType;
import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.RequestBody;
import okhttp3.Response;
import okhttp3.ResponseBody;
import org.springframework.stereotype.Component;
import top.ochiamalu.aioj.properties.AIProperties;
import top.ochiamalu.aioj.util.model.AIResponse;
import top.ochiamalu.aioj.util.model.AIResult;
import top.ochiamalu.aioj.util.model.AIStreamCallback;
import top.ochiamalu.aioj.util.model.Delta;
import top.ochiamalu.aioj.util.model.Message;
import top.ochiamalu.aioj.util.model.StreamChoice;
import top.ochiamalu.aioj.util.model.StreamResponse;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

/**
 * AI工具类
 *
 * @author ochiamalu
 * @date 2025/03/03
 */
@Component
@Slf4j
public class AIUtils {

    private static AIProperties aiProperties;

    @Resource
    private AIProperties tempProperties;

    /**
     * URL
     */
    private static final String URL = "https://api.deepseek.com/chat/completions";
    /**
     * gson
     */
    private static final Gson GSON = new Gson();

    /**
     * 实例
     */
    private static volatile AIUtils instance;
    /**
     * 客户端
     */
    private final OkHttpClient client;

    private AIUtils() {
        this.client = new OkHttpClient.Builder()
                .connectTimeout(5, TimeUnit.MINUTES)
                .writeTimeout(5, TimeUnit.MINUTES)
                .readTimeout(5, TimeUnit.MINUTES)
                .build();
    }

    public static AIUtils getInstance() {
        if (instance == null) {
            synchronized (AIUtils.class) {
                if (instance == null) {
                    instance = new AIUtils();
                }
            }
        }
        return instance;
    }

    /**
     * 请求体构建器
     */
    private static class RequestBodyBuilder {
        private final JsonObject requestObj;
        private final JsonArray messages;

        private RequestBodyBuilder(boolean useR1) {
            this.requestObj = new JsonObject();
            this.messages = new JsonArray();
            requestObj.addProperty("model", useR1 ? "deepseek-reasoner" : "deepseek-chat");
        }

        public RequestBodyBuilder addMessage(String content) {
            JsonObject messageObj = new JsonObject();
            messageObj.addProperty("role", "user");
            messageObj.addProperty("content", content);
            messages.add(messageObj);
            return this;
        }

        public RequestBodyBuilder setStream(boolean useStream) {
            requestObj.addProperty("stream", useStream);
            return this;
        }

        public String build() {
            requestObj.add("messages", messages);
            return GSON.toJson(requestObj);
        }
    }

    /**
     * 消息发送策略接口
     */
    private interface MessageSendStrategy {
        AIResult send(boolean useStream);
    }

    /**
     * 单消息发送策略
     */
    private class SingleMessageStrategy implements MessageSendStrategy {
        private final String message;
        private final boolean useR1;

        public SingleMessageStrategy(String message, boolean useR1) {
            this.message = message;
            this.useR1 = useR1;
        }

        @Override
        public AIResult send(boolean useStream) {
            String requestBody = new RequestBodyBuilder(useR1)
                    .addMessage(message)
                    .setStream(useStream)
                    .build();
            return executeRequest(requestBody);
        }
    }

    /**
     * 多消息发送策略
     */
    private class MultiMessageStrategy implements MessageSendStrategy {
        private final List<String> messages;
        private final boolean useR1;

        public MultiMessageStrategy(List<String> messages, boolean useR1) {
            this.messages = messages;
            this.useR1 = useR1;
        }

        @Override
        public AIResult send(boolean useStream) {
            RequestBodyBuilder builder = new RequestBodyBuilder(useR1);
            messages.forEach(builder::addMessage);
            String requestBody = builder.setStream(useStream).build();
            return executeRequest(requestBody);
        }
    }

    /**
     * 执行HTTP请求
     */
    private AIResult executeRequest(String requestBody) {
        if (!aiProperties.isEnable()) {
            return new AIResult("AI未启用", "");
        }
        MediaType JSON = MediaType.parse("application/json; charset=utf-8");
        RequestBody body = RequestBody.create(requestBody, JSON);

        Request request = new Request.Builder()
                .url(URL)
                .addHeader("Content-Type", "application/json")
                .addHeader("Authorization", "Bearer " + aiProperties.getApiKey())
                .post(body)
                .build();

        try (Response response = client.newCall(request).execute()) {
            if (response.body() != null) {
                String responseBody = response.body().string();
                AIResponse aiResponse = GSON.fromJson(responseBody, AIResponse.class);
                Message message = aiResponse.getChoices()[0].getMessage();
                return new AIResult(message.getContent(), message.getReasoning_content());
            }
        } catch (Exception e) {
            log.error("发生错误: {}", e.getMessage(), e);
        }
        return null;
    }

    /**
     * 发送单条消息
     */
    public AIResult sendMessage(String message, boolean useStream, boolean useR1) {
        return new SingleMessageStrategy(message, useR1).send(useStream);
    }

    /**
     * 发送多条消息
     */
    public AIResult sendMessageList(List<String> messageList, boolean useStream, boolean useR1) {
        return new MultiMessageStrategy(messageList, useR1).send(useStream);
    }

    /**
     * 执行流式HTTP请求
     */
    private void executeStreamRequest(String requestBody, AIStreamCallback callback) {
        MediaType JSON = MediaType.parse("application/json; charset=utf-8");
        RequestBody body = RequestBody.create(requestBody, JSON);

        Request request = new Request.Builder()
                .url(URL)
                .addHeader("Content-Type", "application/json")
                .addHeader("Authorization", "Bearer " + aiProperties.getApiKey())
                .post(body)
                .build();

        CountDownLatch latch = new CountDownLatch(1);

        client.newCall(request).enqueue(new Callback() {
            @Override
            public void onFailure(Call call, IOException e) {
                callback.onError(e.getMessage());
                latch.countDown();
            }

            @Override
            public void onResponse(Call call, Response response) {
                try (ResponseBody responseBody = response.body()) {
                    if (responseBody == null) {
                        callback.onError("Response body is null");
                        latch.countDown();
                        return;
                    }

                    BufferedReader reader = new BufferedReader(
                            new InputStreamReader(responseBody.byteStream(), StandardCharsets.UTF_8));
                    String line;
                    while ((line = reader.readLine()) != null) {
                        if (line.startsWith("data: ")) {
                            String jsonData = line.substring(6);
                            if ("[DONE]".equals(jsonData)) {
                                callback.onComplete();
                                latch.countDown();
                                continue;
                            }

                            try {
                                StreamResponse streamResponse = GSON.fromJson(jsonData, StreamResponse.class);
                                if (streamResponse.getChoices() != null &&
                                        streamResponse.getChoices().length > 0) {

                                    StreamChoice choice = streamResponse.getChoices()[0];
                                    if (choice.getDelta() != null) {
                                        Delta delta = choice.getDelta();
                                        callback.onMessage(
                                                delta.getContent(),
                                                delta.getReasoning_content()
                                        );
                                    }
                                }
                            } catch (JsonSyntaxException e) {
                                log.error("JSON解析错误: {}", e.getMessage(), e);
                            }
                        }
                    }
                } catch (Exception e) {
                    callback.onError(e.getMessage());
                    latch.countDown();
                }
            }
        });

        try {
            if (!latch.await(5, TimeUnit.MINUTES)) {
                callback.onError("操作超时");
            }
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            callback.onError("操作被中断");
        }
    }

    /**
     * 发送流式单条消息
     */
    public void sendStreamMessage(String message, boolean useR1, AIStreamCallback callback) {
        String requestBody = new RequestBodyBuilder(useR1)
                .addMessage(message)
                .setStream(true)
                .build();
        executeStreamRequest(requestBody, callback);
    }

    /**
     * 发送流式多条消息
     */
    public void sendStreamMessageList(List<String> messageList, boolean useR1, AIStreamCallback callback) {
        RequestBodyBuilder builder = new RequestBodyBuilder(useR1);
        messageList.forEach(builder::addMessage);
        String requestBody = builder.setStream(true).build();
        executeStreamRequest(requestBody, callback);
    }

    @PostConstruct
    public void init() {
        aiProperties = tempProperties;
    }
}
